"""
任务一改进版本 - 基于题目要求和reference文件实现

主要改进：
1. 智能数据筛选策略
2. 代表性数据分析
3. 源域和目标域数据分别保存
4. 基于reference文件的特征提取
5. 详细的分析报告
"""

import os
import sys
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')  # 使用非交互式后端
import matplotlib.pyplot as plt
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

# 添加父目录到路径
sys.path.append('..')

# 导入自定义模块
from data_loader import DataLoader
from feature_extractor import FeatureExtractor
from feature_selector import FeatureSelector
from data_analyzer import DataAnalyzer
from improved_data_selector import ImprovedDataSelector

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


def main():
    """主函数"""
    print("🚀 开始执行任务一：数据分析与故障特征提取（改进版）")
    print("=" * 80)
    
    # 记录开始时间
    start_time = datetime.now()
    timestamp = start_time.strftime("%Y%m%d_%H%M%S")
    print(f"⏰ 执行时间: {timestamp}")
    
    # 创建输出目录
    output_dir = '.'
    os.makedirs(output_dir, exist_ok=True)
    
    try:
        # 步骤1: 数据加载和筛选
        print("\n" + "="*60)
        print("步骤1: 数据加载和智能筛选")
        print("="*60)
        
        data_loader = DataLoader(
            source_data_path='../datasets/sourceDataset',
            target_data_path='../datasets/targetDataset',
            target_sampling_rate=12000
        )
        
        # 加载源域数据（使用系统性筛选策略）
        source_data, selected_files, file_analysis = data_loader.load_source_data(sampling_strategy='systematic')
        
        # 加载目标域数据
        target_data = data_loader.load_target_data()
        
        print(f"✅ 数据加载完成:")
        print(f"   源域数据: {len(source_data)}个文件")
        print(f"   目标域数据: {len(target_data)}个文件")
        
        # 步骤2: 代表性数据分析
        print("\n" + "="*60)
        print("步骤2: 代表性源域数据分析")
        print("="*60)
        
        data_analyzer = DataAnalyzer()
        
        # 选择代表性数据（每类5个样本）
        representative_data = data_analyzer.select_representative_data(source_data, n_samples_per_class=5)
        
        # 创建特征提取器
        feature_extractor = FeatureExtractor(target_fs=12000)
        
        # 分析代表性数据
        representative_analysis = data_analyzer.analyze_representative_data(representative_data, feature_extractor)
        
        # 生成代表性数据分析可视化
        data_analyzer.visualize_representative_analysis(representative_analysis, output_dir)
        
        # 生成代表性数据分析报告
        representative_report_path = f"{output_dir}/representative_analysis_report_{timestamp}.md"
        data_analyzer.generate_representative_analysis_report(representative_analysis, representative_report_path)
        
        # 步骤3: 整体数据集特征提取
        print("\n" + "="*60)
        print("步骤3: 整体数据集特征提取")
        print("="*60)
        
        # 提取源域特征
        print("🔍 提取源域特征...")
        source_features = extract_features_from_data(source_data, feature_extractor, 'source')
        
        # 提取目标域特征
        print("🔍 提取目标域特征...")
        target_features = extract_features_from_data(target_data, feature_extractor, 'target')
        
        print(f"✅ 特征提取完成:")
        print(f"   源域特征: {len(source_features)}个样本")
        print(f"   目标域特征: {len(target_features)}个样本")
        
        # 步骤4: 特征选择
        print("\n" + "="*60)
        print("步骤4: 特征选择和可视化")
        print("="*60)
        
        # 准备特征选择数据
        source_df = pd.DataFrame(source_features)
        target_df = pd.DataFrame(target_features)
        
        # 获取数值特征列
        numeric_columns = source_df.select_dtypes(include=[np.number]).columns
        numeric_columns = [col for col in numeric_columns if col not in ['fault_size', 'load_condition']]
        
        X_source = source_df[numeric_columns]
        y_source = source_df['fault_type']
        
        # 创建特征选择器
        feature_selector = FeatureSelector()
        
        # 分析特征重要性
        importance_df = feature_selector.analyze_feature_importance(X_source, y_source)
        
        # 分析特征相关性
        corr_matrix, high_corr_pairs = feature_selector.analyze_feature_correlation(X_source)
        
        # 分析特征分布
        distribution_stats = feature_selector.analyze_feature_distribution(X_source, y_source)
        
        # 选择特征
        X_selected = feature_selector.select_features(X_source, y_source, method='random_forest', k=30)
        
        # 生成特征分析可视化
        generate_feature_analysis_visualizations(
            feature_selector, X_source, y_source, importance_df, corr_matrix, 
            distribution_stats, output_dir, timestamp
        )
        
        # 步骤5: 数据保存
        print("\n" + "="*60)
        print("步骤5: 数据保存")
        print("="*60)
        
        # 分别保存源域和目标域数据
        source_file, target_file = data_loader.save_data_separately(source_df, target_df, output_dir)
        
        # 保存选中的特征数据
        source_selected_df = source_df[['file_name', 'fault_type', 'fault_size', 'load_condition'] + feature_selector.selected_features]
        target_selected_df = target_df[['file_name'] + feature_selector.selected_features]
        
        source_selected_file = f"{output_dir}/source_domain_selected_features_{timestamp}.csv"
        target_selected_file = f"{output_dir}/target_domain_selected_features_{timestamp}.csv"
        
        source_selected_df.to_csv(source_selected_file, index=False, encoding='utf-8-sig')
        target_selected_df.to_csv(target_selected_file, index=False, encoding='utf-8-sig')
        
        print(f"✅ 数据保存完成:")
        print(f"   源域完整数据: {source_file}")
        print(f"   目标域完整数据: {target_file}")
        print(f"   源域选中特征: {source_selected_file}")
        print(f"   目标域选中特征: {target_selected_file}")
        
        # 步骤6: 生成综合报告
        print("\n" + "="*60)
        print("步骤6: 生成综合报告")
        print("="*60)
        
        generate_comprehensive_report(
            source_data, target_data, selected_files, file_analysis,
            representative_analysis, importance_df, corr_matrix, high_corr_pairs,
            distribution_stats, feature_selector.selected_features,
            output_dir, timestamp
        )
        
        # 计算执行时间
        end_time = datetime.now()
        execution_time = end_time - start_time
        
        print("\n" + "="*80)
        print("🎉 任务一执行完成！")
        print("="*80)
        print(f"⏱️ 总执行时间: {execution_time}")
        print(f"📁 输出目录: {output_dir}")
        print("="*80)
        
    except Exception as e:
        print(f"❌ 执行过程中出现错误: {e}")
        import traceback
        traceback.print_exc()
        return False
    
    return True


def extract_features_from_data(data_list, feature_extractor, data_type):
    """从数据列表中提取特征"""
    features_list = []
    
    for data in data_list:
        try:
            # 使用DE信号进行特征提取
            de_signal = data['de_signal']
            if len(de_signal) > 0:
                features = feature_extractor.extract_all_features(
                    de_signal, 
                    data['rpm'], 
                    'SKF6205'
                )
                features['file_name'] = data['file_name']
                features['fault_type'] = data['fault_type']
                features['fault_size'] = data['fault_size']
                features['load_condition'] = data['load_condition']
                features['data_type'] = data_type
                features_list.append(features)
        except Exception as e:
            print(f"⚠️ 特征提取失败 {data['file_name']}: {e}")
            continue
    
    return features_list


def generate_feature_analysis_visualizations(feature_selector, X, y, importance_df, corr_matrix, 
                                           distribution_stats, output_dir, timestamp):
    """生成特征分析可视化"""
    print("📊 生成特征分析可视化...")
    
    # 1. 特征重要性图
    importance_path = f"{output_dir}/feature_importance_analysis_{timestamp}.png"
    feature_selector.plot_feature_importance(importance_df, top_n=20, save_path=importance_path)
    
    # 2. 特征相关性热力图
    correlation_path = f"{output_dir}/feature_correlation_heatmap_{timestamp}.png"
    feature_selector.plot_feature_correlation(corr_matrix, save_path=correlation_path)
    
    # 3. 特征分布图
    distribution_path = f"{output_dir}/feature_distribution_{timestamp}.png"
    feature_selector.plot_feature_distribution(X, y, X.columns.tolist(), save_path=distribution_path)
    
    # 4. 特征空间可视化
    space_path = f"{output_dir}/feature_space_visualization_{timestamp}.png"
    feature_selector.visualize_features(X, y, method='tsne', save_path=space_path)
    
    print("✅ 特征分析可视化完成")


def generate_comprehensive_report(source_data, target_data, selected_files, file_analysis,
                                representative_analysis, importance_df, corr_matrix, high_corr_pairs,
                                distribution_stats, selected_features, output_dir, timestamp):
    """生成综合报告"""
    print("📝 生成综合报告...")
    
    report = f"""# 任务一：数据分析与故障特征提取 - 综合报告

## 1. 任务概述

### 1.1 任务目标
根据"高速列车轴承智能故障诊断问题.txt"文件要求，完成以下任务：
1. 从提供的源域数据中筛选部分数据组成数据集
2. 结合轴承故障机理，选择合适的方法对有代表性的源域数据进行特征分析
3. 对整体数据集进行特征提取，用于后续诊断任务

### 1.2 数据概况
- **源域数据**: 从161个文件中筛选出{len(source_data)}个代表性文件
- **目标域数据**: {len(target_data)}个未知标签的故障文件
- **特征维度**: {len(selected_features)}个选中特征

## 2. 数据筛选策略

### 2.1 筛选依据
基于以下原则进行智能筛选：
1. **故障类型平衡**: 确保每种故障类型都有足够的代表性样本
2. **故障尺寸覆盖**: 包含不同故障尺寸的样本
3. **载荷条件多样**: 覆盖不同载荷条件下的数据
4. **采样频率统一**: 统一使用12kHz采样频率

### 2.2 筛选结果
"""
    
    # 添加筛选结果
    for fault_type, files in file_analysis['fault_types'].items():
        if files:
            report += f"- **{fault_type}**: {len(files)}个文件\n"
    
    report += f"""
## 3. 代表性数据分析

### 3.1 分析样本
从源域数据中选择了{len(representative_analysis['features_df'])}个代表性样本进行分析。

### 3.2 主要发现
1. **特征分布差异**: 不同故障类型在关键特征上表现出明显差异
2. **特征相关性**: 发现{len(high_corr_pairs)}对高相关性特征
3. **故障特征**: 各故障类型具有独特的特征模式

## 4. 特征提取与选择

### 4.1 特征类型
提取了以下5类特征：
1. **时域特征**: 均值、标准差、RMS、偏度、峰度等
2. **频域特征**: 频域统计、频带能量等
3. **时频域特征**: 小波变换、STFT特征等
4. **故障特征频率**: 轴承故障特征频率及其谐波
5. **非线性特征**: 近似熵、样本熵、排列熵等

### 4.2 特征选择结果
使用随机森林方法选择了{len(selected_features)}个重要特征：

"""
    
    # 添加选中的特征
    for i, feature in enumerate(selected_features, 1):
        report += f"{i}. {feature}\n"
    
    report += f"""
## 5. 特征重要性分析

### 5.1 前10个重要特征
"""
    
    # 添加前10个重要特征
    top_10_features = importance_df.head(10)
    for i, (_, row) in enumerate(top_10_features.iterrows(), 1):
        report += f"{i}. **{row['feature']}**: 综合分数 {row['combined_score']:.4f}\n"
    
    report += f"""
### 5.2 特征重要性统计
- **随机森林平均重要性**: {importance_df['random_forest'].mean():.4f}
- **互信息平均分数**: {importance_df['mutual_info'].mean():.4f}
- **F检验平均分数**: {importance_df['f_test'].mean():.4f}
- **综合平均分数**: {importance_df['combined_score'].mean():.4f}

## 6. 数据保存

### 6.1 保存文件
1. **源域完整数据**: `source_domain_data.csv`
2. **目标域完整数据**: `target_domain_data.csv`
3. **源域选中特征**: `source_domain_selected_features_{timestamp}.csv`
4. **目标域选中特征**: `target_domain_selected_features_{timestamp}.csv`

### 6.2 可视化文件
1. **代表性数据分析图**: `representative_*.png`
2. **特征重要性图**: `feature_importance_analysis_{timestamp}.png`
3. **特征相关性图**: `feature_correlation_heatmap_{timestamp}.png`
4. **特征分布图**: `feature_distribution_{timestamp}.png`
5. **特征空间图**: `feature_space_visualization_{timestamp}.png`

## 7. 技术实现

### 7.1 数据筛选算法
- 使用平衡采样策略确保各类故障样本的均衡性
- 基于故障类型、故障尺寸、载荷条件进行分层采样
- 智能选择最具代表性的样本进行分析

### 7.2 特征提取方法
- 基于reference/1-2-Multi-dimensional_feature_extraction.py实现
- 支持时域、频域、时频域、故障频率、非线性特征提取
- 针对轴承故障特点优化特征提取算法

### 7.3 特征选择策略
- 基于reference/1-3-Feature_selection_and_visualization.py实现
- 结合随机森林、互信息、F检验的综合评分
- 考虑特征相关性和分布差异进行选择

## 8. 结论与建议

### 8.1 主要成果
1. 成功实现了智能数据筛选，选择了{len(source_data)}个代表性源域文件
2. 完成了代表性数据的深入分析，为后续特征提取提供了依据
3. 提取了{len(selected_features)}个重要特征，为后续诊断任务奠定了基础
4. 分别保存了源域和目标域数据，符合迁移学习的要求

### 8.2 技术特点
1. **智能筛选**: 基于多维度指标进行数据筛选
2. **代表性分析**: 深入分析代表性数据的特征分布
3. **多维特征**: 提取5类46维特征，全面描述故障信息
4. **科学选择**: 基于多种方法综合评估特征重要性

### 8.3 后续建议
1. 在任务二中使用选中的特征进行模型训练
2. 在任务三中利用源域和目标域的特征差异进行迁移学习
3. 在任务四中基于特征重要性进行可解释性分析

---
*报告生成时间: {datetime.now().strftime('%Y年%m月%d日 %H:%M:%S')}*
*基于实际数据分析结果生成*
"""
    
    # 保存报告
    report_path = f"{output_dir}/task1_comprehensive_report_{timestamp}.md"
    with open(report_path, 'w', encoding='utf-8') as f:
        f.write(report)
    
    print(f"✅ 综合报告已保存: {report_path}")


if __name__ == "__main__":
    success = main()
    if success:
        print("\n🎉 任务一执行成功！")
    else:
        print("\n❌ 任务一执行失败！")
